#!/usr/bin/env python  
#-*- coding:utf-8 _*-  
""" 
@author:hello_life 
@license: Apache Licence 
@file: train_utils.py 
@time: 2022/05/02
@software: PyCharm 
description:
"""
import os

import torch
import torch.nn as nn

from utils.data_utils import get_f1,id2label

def train_loop(model,train_dataloader,test_dataloader,config):
    optimzier=torch.optim.Adam(model.parameters())
    for epoch in range(config.epochs):
        model.train()
        total_loss=0.0
        y_pred = []
        y_label = []
        for i,(x,mask,y) in enumerate(train_dataloader):
            tag_scores,loss=model.forward_with_crf(x,mask,y)
            preds=model.crf.decode(tag_scores,mask)
            loss.backward()

            #梯度裁剪
            # nn.utils.clip_grad_norm(model.parameters(),config.clip_lr)
            optimzier.step()
            optimzier.zero_grad()

            total_loss+=loss.item()
            for j, p in enumerate(preds):
                y_pred.append(preds[j])
                y_label.append(y[j].cpu().tolist())
            if i>0 and i%config.steps==0:
                #模型存储
                if not os.path.exists(config.save_dir):
                    os.makedirs(config.save_dir)
                torch.save({
                    "epoch":epoch,
                    "step":i,
                    "model_state_dict":model.state_dict(),
                    "optimizer_state_dict":optimzier.state_dict(),
                },config.save_path)

                f1_score_total, f1_score_class = get_f1(y_pred, y_label, config)
                y_pred,y_label=[],[]
                print(f"Train_Process: epoch:{epoch},steps:{i},loss:{loss.item()},"
                      f"f1_score_total:{f1_score_total},f1_score_class:{f1_score_class}")
        test_loop(model,test_dataloader, config)


def test_loop(model,dataloader,config):
    model.eval()
    y_pred=[]
    y_label=[]
    total_loss=0.0
    with torch.no_grad():
        for i,(x,mask,y) in enumerate(dataloader):
            output,loss = model.forward_with_crf(x,mask,y)
            pred = model.crf.decode(output,mask)

            total_loss+=loss.item()

            #因为y_pred形状是[[[1],[2]],[[2],[3]]],转化为 [[1],[2],[3],[4]]
            for i,p in enumerate(pred):
                y_pred.append(pred[i])
                y_label.append(y[i].cpu().tolist())

    f1_score_total,f1_score_class=get_f1(y_pred,y_label,config)
    avg_loss=total_loss/len(dataloader)
    print(f"Test_process，loss:{avg_loss},f1_score_total:{f1_score_total},"
          f"f1_score_class:{f1_score_class}")